import numpy as np
import pandas as pd
import talib as ta

from datetime import datetime
from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.database import database_manager
C0LOR1 = "PapayaWhip"
COLOR2 = "Tan"
#functionS
def my_func(x):
    return [x+1, x+2]

def my_func(x):
    return [x+1, x+2]
selectColor = lambda num: C0LOR1 if num % 2 == 0 else COLOR2
# 使用APPLY


#MYFUNCTION
def test_lambda():
    for i in range(10):
        COLOR = selectColor(i) # "Tan
        print(COLOR)
def test_dataframe():
    df = pd.DataFrame({'x':[1,2,3,4,5]})
    df['y'], df['z'] = df.apply(lambda row: my_func(row.x), axis = 1)# failed to apply lambda to myfunction
    print(df)
    df['y'], df['z'] = my_func(df.x)#success
    print(df)
    f = lambda x : x.max() - x.min() 
    # test row and col
    df = pd.DataFrame(np.random.randn(4, 3), columns=list('bde'), index=['utah', 'ohio', 'texas', 'oregon']) #columns表述列标， index表述行标
    print(df)

    t1 = df.apply(f) #df.apply(function, axis=0)，默认axis=0，表示将一列数据作为Series的数据结构传入给定的function中
    print(t1)

    t2 = df.apply(f, axis=1)
    print(t2)

if __name__ == '__main__':
    # df = pd.read_csv('IF2201_CFFEX_3056.csv', header=None)
    symbol = "IF2201"
    exchange = Exchange.CFFEX
    start = datetime(2021, 12, 1)
    end = datetime(2022, 3, 17)
    interval = Interval.MINUTE
    bars = database_manager.load_bar_data(
        symbol, Exchange.CFFEX, interval=Interval.MINUTE, start=start, end=end
    )
    n =30
    history = bars[:n]
    left = len(bars) - n
    new_data = bars[n:]

    df = pd.DataFrame(
        [
            {
                "Close": s.close_price,
                "Low": s.low_price,
                "High": s.high_price,
                "Open": s.open_price,
                "Volume": s.volume,
            }
            for s in history
        ]
    )

    dd = None
    print(datetime.now())

    for i in range(len(new_data)):
        s = new_data.pop(0)
        df = df.append(
            {
                "Close": s.close_price,
                "Low": s.low_price,
                "High": s.high_price,
                "Open": s.open_price,
                "Volume": s.volume,
            },
            ignore_index=True,
        )
        df["Mean"] = ta.SMA(df["Close"], n)
    print(df.tail(3))
    print(datetime.now())

    df2 = pd.DataFrame(
        [
            {
                "Close": s.close_price,
                "Low": s.low_price,
                "High": s.high_price,
                "Open": s.open_price,
                "Volume": s.volume,
                'mean':np.nan
            }
            for s in history
        ]
    )
    new_data = new_data = bars[n:]
    mean = np.mean(df["Close"].tail(n))
    for i in range(len(new_data)):
        s = new_data.pop(0)
        
        mean = mean + ( s.close_price - df.iloc[-n:].Close )/n
        # print(df.iloc[-n+1:],mean)
        df2 = df2.append(
            {
                "Close": s.close_price,
                "Low": s.low_price,
                "High": s.high_price,
                "Open": s.open_price,
                "Volume": s.volume,
                'mean':mean
            },
            ignore_index=True,
        )
        # avg = np.average(df["Close"].head(n))
        # mean = np.mean(df2["Close"].tail(n))
        # df["Mean2"] = np.mean(df["Close"].tail(30))
    print(df2.tail(3))
    print(datetime.now())
    for i in range(30):

        ta.SMA(float(df["Close"]),30)
    print(datetime.now())
    for i in range(30):
        np.mean(float(df[0:30]['Close']))
    print(datetime.now())
    

